# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: chensili

import logging
import os
from typing import List, Optional, Union

from coin.base.config import Config
from coin.base.param_util import to_list
from coin.base.timestamp import get_timestamp

from coin.exchange.base.fill_stats import FillStatsCalculator2
from coin.exchange.base.order_gateway_logger import gen_og_log_request

from coin.exchange.base.order_gateway import (OrderGatewayBase,
                                              Order,
                                              OrderSide,
                                              OrderType,
                                              OrderGatewayStatus)
from coin.exchange.base.order_fill_util import FillManager

from coin.exchange.bitfinex_v2.kr_rest.product import (
    BitfinexCurrency,
    BitfinexProduct,
)
from coin.exchange.bitfinex_v2.og_util import BitfinexOrderInternal
from coin.exchange.bitfinex_v2.order_gateway_log import (
    BitfinexV2OrderGatewayLogProvider,)
from coin.exchange.bitfinex_v2.ws.private_subscriber import (BitfinexV2PrivateSubscriber)
from coin.exchange.base.order_gateway import BaseOrderGatewayInfoMutator
import coin.flow.subscriber as flow_subscriber

from coin.proto.coin_order_gateway_pb2 import OrderEvent
from coin.proto.coin_query_pb2 import (CurrencyBalance,
                                       AccountBalance,
                                       AccountPosition,
                                       ProductPosition)


def generate_client_order_id():
  timestamp = int(get_timestamp() / 1000)
  return int(timestamp % 10**12)


class BitfinexV2FillStatsCalculator(FillStatsCalculator2):
  ProductType = BitfinexProduct
  BuyOrderSide = [OrderSide.BUY]
  SellOrderSide = [OrderSide.SELL]

  def __init__(self, products):
    super().__init__(products)


class BitfinexV20GatewayInfo(BaseOrderGatewayInfoMutator):
  CurrencyType = BitfinexCurrency
  ProductType = BitfinexProduct
  BalanceType = float
  PositionType = float

  def __init__(self, currencies, products):
    super().__init__(to_list(currencies), to_list(products))

    self._funding_balances = {currency.symbol: 0. for currency in currencies}

    self.set_zero_balances()
    self.set_zero_positions()

    # TODO(inkyu): Make it as OKEX futures OG
    self._avg_entry_prices = {product: 0. for product in self.products}

  def get_average_entry_price(self, product):
    return self._avg_entry_prices.get(product, 0.0)

  def get_funding_balances(self):
    return self._funding_balances

  def get_funding_balance(self, currency):
    if isinstance(currency, self.CurrencyType):
      return self._funding_balances.get(currency.symbol, self.BalanceEmptyValue)
    elif isinstance(currency, str):
      return self._funding_balances[currency]
    else:
      raise TypeError("Not Supported: %s" % (currency,))

  def update_bitfinex_position(self, product, qty, base_price):
    self.set_position(product, qty)
    self._avg_entry_prices[product] = base_price

  def set_funding_balance(self, currency, balance):
    self.assert_type_currency(currency)
    assert isinstance(balance, self.BalanceType)
    self._funding_balances[currency.symbol] = balance
    return balance

  def update_position_by_trade(self, product, position):
    self.assert_type_product(product)
    assert isinstance(position, self.PositionType)
    new_position = self.get_position(product) + position
    self.set_position(product, new_position)

  def update_funding_balance_by_funding_trade(self, currency, balance):
    self.assert_type_currency(currency)
    assert isinstance(balance, self.BalanceType)
    new_balance = self.get_funding_balance(currency) + (-1) * balance
    self.set_funding_balance(currency, new_balance)

  def balance_as_proto(self):
    currency_balance_list = []
    for currency, balance in self.get_balances().items():
      currency_balance = CurrencyBalance(currency=currency, total=balance)
      currency_balance_list.append(currency_balance)
    return AccountBalance(exchange='Bitfinex', each_balance=currency_balance_list)

  def position_as_proto(self):
    position_list = []
    for symbol, position in self.get_positions().items():
      product_position = ProductPosition(symbol=symbol, net_position=position)
      position_list.append(product_position)
    return AccountPosition(exchange='Bitfinex', market_type="Margin", each_position=position_list)

  def funding_balance_as_proto(self):
    currency_balance_list = []
    for currency, balance in self.get_funding_balances().items():
      currency_balance = CurrencyBalance(currency=currency, total=balance)
      currency_balance_list.append(currency_balance)
    return AccountBalance(exchange='Bitfinex',
                          market_type="Margin",
                          each_balance=currency_balance_list)


class BitfinexV2OrderGatewayMargin(OrderGatewayBase):
  _exchange = 'Bitfinex'
  _market_type = 'Margin'

  def __init__(
      self,
      currencies: Union[List[BitfinexCurrency], BitfinexCurrency],
      products: Union[List[BitfinexProduct], BitfinexProduct],
      config: Config,  # Config(filename='bitfinex')
      logger=None):
    super().__init__()
    self._logger = logger or logging.getLogger(__name__)

    self._config = config

    self._working_orders = {}
    self._order = None

    self._account_ready = False
    self._orders_ready = False
    self._og_info = BitfinexV20GatewayInfo(to_list(currencies), to_list(products))
    self._og_logger = BitfinexV2OrderGatewayLogProvider(gen_og_log_request(config, self))

    self._fill_stats_calculator = BitfinexV2FillStatsCalculator(products)

  def get_fill_stats(self, product):
    return self._fill_stats_calculator.get_fill_stats(product)

  def get_average_entry_price(self, product):
    return self._og_info.get_average_entry_price(product)

  def on_private_sub_closed(self):
    self._og_ready.set_status(OrderGatewayStatus.DISCONNECTED)
    self._logger.error('Exit strategy due to private subscriber closed.')
    os._exit(1)

  def start(self, queue, ioloop, flow_sub=None):
    self._ioloop = ioloop
    self._private_sub = BitfinexV2PrivateSubscriber(
        self._config, queue, 1, close_callback=self.on_private_sub_closed)  # TODO(inkyu): WORKER_ID
    self._private_sub.start(ioloop)

    self._flow_sub = flow_sub or flow_subscriber.from_queue(queue)
    self._flow_sub.subscribe('bitfinex_v2_ws_private_1', self.on_private_msg)

  def stop(self):
    self._flow_sub.unsubscribe('bitfinex_v2_ws_private_1', self.on_private_msg)

  def _check_order_scope(self, native_symbol, order_type):
    ret = True
    # filter out orders of unsubscribed products and different market
    # type
    product = BitfinexProduct.FromStrNativeProduct(native_symbol)
    if product not in self._og_info.products or order_type != 'LIMIT':
      ret = False
    return ret

  def _handle_ws_position_snapshot(self, messages: list):
    for data in messages:
      native_symbol = data[0]
      status = data[1]
      qty = float(data[2])
      base_price = float(data[3])  # Avg entry price
      # unrealized_pnl = data[6]  # None
      # price_liq
      if status == 'ACTIVE':
        product = BitfinexProduct.FromStrNativeProductNothrow(native_symbol)
        if product is None:
          continue
        self._og_info.update_bitfinex_position(product, qty, base_price)
        self._og_logger.write_position(self._og_info.position_as_proto())

  def _handle_ws_wallet(self, message: list):
    wallet_type = message[0]
    native_currency = message[1]
    balance = float(message[2])
    # unsettled_interest = message[3]
    currency = BitfinexCurrency.FromStrNativeCurrencyNothrow(native_currency)
    if currency is None:
      return
    if wallet_type == 'margin':  # exchange, margin, funding
      self._og_info.set_balance(currency, balance)

  def _handle_ws_wallets(self, messages: list):
    for wallet in messages:
      self._handle_ws_wallet(wallet)
    self._og_logger.write_balance(self._og_info.balance_as_proto())
    self._account_ready = True
    self._check_ready()

  def _handle_ws_order_update(self, given_order: list, channel_id):
    self._logger.debug('%s: %s', channel_id, given_order)
    server_order_id = given_order[0]
    group_id = given_order[1]
    client_order_id = given_order[2]
    symbol = given_order[3]
    # mts_create = given_order[4]
    # mts_update = given_order[5]
    qty = given_order[6]  # remaining qty
    qty_orig = given_order[7]
    order_type = given_order[8]  # LIMIT, MARKET, STOP, TRAILING STOP, EXCHANGE MARKET,
    order_type_prev = given_order[9]  # EXCHANGE LIMIT, EXCHANGE STOP, EXCHANGE TRAILING STOP,
    flags = given_order[12]
    order_status = given_order[13]  # ACTIVE, EXECUTED, PARTIALLY FILLED, CANCELED
    price_orig = given_order[16]
    price_avg = given_order[17]
    price_trailing = given_order[18]
    price_aux_limit = given_order[19]

    if not self._check_order_scope(symbol, order_type):
      return

    order_side = None
    if qty_orig > 0:
      order_side = OrderSide.BUY
    elif qty_orig < 0:
      order_side = OrderSide.SELL
    else:
      raise ValueError('Zero qty: %s' % str(given_order))

    timestamp = get_timestamp()
    product = BitfinexProduct.FromStrNativeProduct(symbol)
    client_order_id = client_order_id or generate_client_order_id()
    foreign = client_order_id not in self._working_orders
    if client_order_id in self._working_orders:
      order = self._working_orders[client_order_id]
    else:
      order = Order(order_id=client_order_id,
                    product=product,
                    price=price_orig,
                    qty=abs(qty_orig),
                    order_side=order_side,
                    order_type=OrderType.LIMIT,
                    tag=None,
                    accepted=True,
                    cancel_sent=False,
                    internal=BitfinexOrderInternal.create_by_order_update(
                        price_orig=price_orig,
                        qty_orig=abs(qty_orig),
                        server_order_id=server_order_id,
                        group_id=group_id,
                        timestamp_created=timestamp,
                        foreign=foreign))
    order.accepted = True
    order.internal.server_order_id = server_order_id
    order.internal.group_id = group_id

    if foreign:
      self._working_orders[client_order_id] = order
      self._logger.error('Untrackable order: %s', given_order)

    if order_status.startswith('CANCELED') or \
       order_status.startswith('POSTONLY'):
      # Three types of CANCELED status.
      # (1). CANCELED
      # (2). CANCELED was: PARTIALLY FILLED @ PRICE(AMOUNT)
      # (3). POSTONLY CANCELED
      del self._working_orders[client_order_id]
      self._og_logger.gen_order_event_and_log(order=order,
                                              timestamp=timestamp,
                                              order_event_type=OrderEvent.CANCEL_CONFIRMED)
    elif order_status.startswith('EXECUTED @'):
      order.internal.fill.add_fill(price_avg, abs(qty))
      filled_qty = order.internal.fill.value.qty_last_fill
      self._fill_stats_calculator.add_fill_by_order(product=product,
                                                    filled_qty=abs(filled_qty),
                                                    order_side=order_side,
                                                    order_qty=qty_orig)
      self._og_logger.gen_order_event_and_log(order=order,
                                              timestamp=timestamp,
                                              order_event_type=OrderEvent.ORDER_FILLED,
                                              filled_qty=filled_qty)
      del self._working_orders[client_order_id]
    elif order_status.startswith('INSUFFICIENT'):
      # INSUFFICIENT MARGIN was: PARTIALLY FILLED @ PRICE(AMOUNT)
      # INSUFFICIENT BALANCE was: PARTIALLY FILLED @ PRICE(AMOUNT)
      order.internal.fill.add_fill(price_avg, abs(qty))
      filled_qty = order.internal.fill.value.qty_last_fill
      self._fill_stats_calculator.add_fill_by_order(product=product,
                                                    filled_qty=abs(filled_qty),
                                                    order_side=order_side,
                                                    order_qty=qty_orig)
      self._og_logger.gen_order_event_and_log(order=order,
                                              timestamp=timestamp,
                                              order_event_type=OrderEvent.ORDER_FILLED,
                                              filled_qty=filled_qty)
      del self._working_orders[client_order_id]
      self._logger.error('Insufficient margin/balance: %s', given_order)
    elif order_status.startswith('PARTIALLY'):
      order.internal.fill.add_fill(price_avg, abs(qty))
      filled_qty = order.internal.fill.value.qty_last_fill
      self._fill_stats_calculator.add_fill_by_order(product=product,
                                                    filled_qty=abs(filled_qty),
                                                    order_side=order_side,
                                                    order_qty=qty_orig)
      self._og_logger.gen_order_event_and_log(order=order,
                                              timestamp=timestamp,
                                              order_event_type=OrderEvent.ORDER_FILLED,
                                              filled_qty=filled_qty)
    elif order_status == 'ACTIVE':  # Active order
      self._og_logger.gen_order_event_and_log(order=order,
                                              timestamp=timestamp,
                                              order_event_type=OrderEvent.ORDER_ACCEPTED)
    else:
      self._logger.error('Unknown order_status: %s', given_order)

  def _handle_ws_orders_snapshot(self, messages: list):
    channel_id = 'os'
    for given_order in messages:
      self._handle_ws_order_update(given_order, channel_id)
    self._orders_ready = True
    self._check_ready()

  def _handle_ws_order_error(self, message: list):
    payload = message[4]
    status = message[6]
    text = message[7]
    client_order_id = payload[2]
    symbol = payload[3]
    qty = payload[6]
    order_type = payload[8]

    if not self._check_order_scope(symbol, order_type):
      return

    timestamp = get_timestamp()
    if client_order_id in self._working_orders:
      order = self._working_orders[client_order_id]
      self._og_logger.gen_order_event_and_log(order=order,
                                              timestamp=timestamp,
                                              order_event_type=OrderEvent.ORDER_REJECTED)
      del self._working_orders[client_order_id]
    self._logger.error('ORDER %s: %s (%s x %s) %s', status, client_order_id, qty, symbol, text)

  def _handle_ws_trade_event(self, message: list):
    # [225347297, 'tBTCUSD', 1523534381882, 10634024376,
    #  -0.2, 7770.4, None, None, -1]
    trade_id = message[0]
    pair = message[1]
    mts_create = message[2]
    order_id = message[3]
    exec_qty = float(message[4])
    exec_price = float(message[5])
    order_type = message[6]

    if not self._check_order_scope(pair, order_type):
      return

    if message[7] is not None:  # Why is this none?
      order_price = float(message[7])
    else:
      order_price = exec_price
    maker = message[8]

    order_side = None
    if exec_qty > 0:
      order_side = OrderSide.BUY
    elif exec_qty < 0:
      order_side = OrderSide.SELL
    else:
      raise ValueError('Zero qty: "%s"' % str(message))
    product = BitfinexProduct.FromStrNativeProduct(pair)

    self._logger.info('Order filled: %s' % message)
    self._fill_stats_calculator.add_fill_by_trade(product=product,
                                                  filled_qty=abs(exec_qty),
                                                  order_side=order_side)

    timestamp = get_timestamp()
    # TODO(inkyu): Make it log(n)
    order = None
    for elem in self._working_orders.values():
      if elem.internal.server_order_id == order_id:
        order = elem
        break

    if order is None:
      order = Order(order_id=None,
                    product=product,
                    price=order_price,
                    qty=None,
                    order_side=order_side,
                    order_type=OrderType.LIMIT,
                    tag=None,
                    accepted=True,
                    cancel_sent=False,
                    internal=BitfinexOrderInternal.create_by_trade_event(
                        server_order_id=order_id, timestamp_created=timestamp, maker=maker))
    else:
      order.internal.maker = maker
    self._publish_fill(exec_price, abs(exec_qty), order)  # Unsigned

  def _handle_ws_funding_credits_snapshot(self, messages: list):
    for message in messages:
      native_symbol = message[1]
      side = message[2]
      qty = float(message[5])
      status = message[7]
      assert side == -1  # -1 for borrow
      assert qty > 0
      assert status == 'ACTIVE'
      currency = BitfinexCurrency.FromStrNativeCurrencyNothrow(native_symbol)
      if currency is None:
        continue
      self._og_info.set_funding_balance(currency, qty)

  def _handle_ws_funding_trade_event(self, message: list):
    native_symbol = message[1]
    qty = float(message[4])
    currency = BitfinexCurrency.FromStrNativeCurrency(native_symbol)
    assert qty < 0  # Borrow money
    self._og_info.update_funding_balance_by_funding_trade(currency, qty)

  def _handle_ws_funding_credits_close(self, message: list):
    native_symbol = message[1]
    side = message[2]
    qty = float(message[5])
    status = message[7]
    assert qty > 0
    assert side == -1  # -1 for borrow
    assert status.startswith('CLOSED')
    currency = BitfinexCurrency.FromStrNativeCurrency(native_symbol)
    self._og_info.update_funding_balance_by_funding_trade(currency, qty)

  def _handle_ws_notification(self, message: list):
    notify_type = message[1]
    notify_info = message[4]  # Payload
    status = message[6]

    if notify_type == 'on-req':  # About the new order
      if status in ('FAILURE', 'ERROR'):
        self._handle_ws_order_error(message)
      elif status == 'SUCCESS':
        pass
      else:
        self._logger.error('Unknown status: %s (%s)', status, message)
    elif notify_type == 'oc-req':  # About the order cancelation
      if status.startswith('SUCCESS'):
        pass
      elif status.startswith('ERROR'):
        exchange_order_id = notify_info[0]
        found_order = None
        for order_id, order in self._working_orders.items():
          if order.internal.server_order_id == exchange_order_id:
            found_order = order
        if found_order:
          timestamp = get_timestamp()
          self._og_logger.gen_order_event_and_log(order=found_order,
                                                  timestamp=timestamp,
                                                  order_event_type=OrderEvent.CANCEL_ERROR)
          self._logger.error('Cancel error: %s (%s)', status, message)
      else:
        self._logger.error('Unknown cancel status: %s (%s)', status, message)
    elif notify_type == 'oc_multi-req':
      if status == 'INFO':
        pass
      else:
        self._logger.error('Unknown multi-cancel status: %s (%s)', status, message)
    elif notify_type in (
        'uca',
        'fon-req',
        'foc-req',
    ):
      self._logger.error('Not implemented notify_type: %s (%s)', notify_type, message)
    else:
      self._logger.error('Unknown notify_type: %s (%s)', notify_type, message)

  def _handle_ws_position_update(self, data, event_type):
    native_symbol = data[0]
    status = data[1]
    qty = float(data[2])
    base_price = float(data[3])
    if event_type == 'pu':
      assert status == 'ACTIVE', data
    elif event_type == 'pc':
      assert status == 'CLOSED', data
      # set position to 0 when we receive 'pc'
      qty = 0.
    else:
      raise ValueError('Unknown event_type: %s' % event_type)
    product = BitfinexProduct.FromStrNativeProduct(native_symbol)
    self._og_info.update_bitfinex_position(product, qty, base_price)
    self._og_logger.write_position(self._og_info.position_as_proto())

  def on_private_msg(self, record, queue_data, topic_data):
    message = record.data

    event_type = message[1]
    # self._logger.info('%s: (%s)', event_type, message)
    if event_type == 'hb':  # Heartbeat
      # TODO(inkyu): Check heartbeat
      pass
    elif event_type == 'n':  # Notification
      self._handle_ws_notification(message[2])
    # 'os' is only received once when we first log in.
    elif event_type == 'os':
      self._handle_ws_orders_snapshot(message[2])
    # New order, Update order, Cancel order. (Involve fill)
    elif event_type in ('on', 'ou', 'oc'):
      self._handle_ws_order_update(message[2], event_type)
    elif event_type == 'ps':
      self._handle_ws_position_snapshot(message[2])
    # Duplicated position update may be received after order filled.
    elif event_type in ('pu', 'pc'):
      self._handle_ws_position_update(message[2], event_type)
    # 'pn' never received when doing test.
    elif event_type == 'pn':
      # TODO(chensili): delete log.
      self._logger.info(message)
    elif event_type == 'fcs':  # Funding credits snapshot
      self._handle_ws_funding_credits_snapshot(message[2])
    # Ignore 'fos' and 'fls' because margin trading
    # only involves funding credits
    elif event_type in ('fos', 'fls'):  # Funding offers and loans
      pass
    elif event_type == 'fcc':  # Funding credits close
      self._handle_ws_funding_credits_close(message[2])
    elif event_type in ('fcn', 'fcu'):
      pass
    elif event_type == 'fte':
      # Content of 'fte' and 'ftu' is same, so ignore 'fte'
      pass
    elif event_type == 'ftu':  # Funding trade update
      self._handle_ws_funding_trade_event(message[2])
    elif event_type == 'ws':
      self._handle_ws_wallets(message[2])
    elif event_type == 'wu':  # Wallet update
      self._handle_ws_wallet(message[2])
    elif event_type == 'te':
      # Content of 'te' and 'tu' is almost same, so ignore 'tu'
      self._handle_ws_trade_event(message[2])
    elif event_type == 'tu':  # Trade update
      pass
    elif event_type in ('bu', 'miu', 'fiu'):
      # Balance info, Margin info, Funding info
      pass
    else:
      self._logger.error('Unknown event_type: %s (%s)', event_type, message)

  def _check_ready(self):
    if self._account_ready and self._orders_ready:
      self._og_ready.set_status(OrderGatewayStatus.READY)

  def get_working_order(self):
    return self._working_orders.values()

  def has_pending_order(self, product):
    for order in self._working_orders.values():
      if order.product == product and order.internal.server_order_id is None:
        return True
    return False

  def submit_impl(
      self,
      product: BitfinexProduct,
      price: float,
      qty: float,
      order_side: OrderSide,
      order_type: OrderType = OrderType.LIMIT,  # LIMIT
      order_id: Optional[int] = None,
      tag=None,
      post_only=False,
      hidden=False):
    assert isinstance(price, float)
    assert isinstance(qty, float)
    assert isinstance(product, BitfinexProduct)
    assert self.is_ready()

    post_qty = None
    if order_side == OrderSide.BUY:
      post_qty = '%f' % qty  # Positive for buy
    elif order_side == OrderSide.SELL:
      post_qty = '-%f' % qty  # Negative for sell
    else:
      raise ValueError('Unknown side: %s' % order_side)
    if order_type != OrderType.LIMIT:
      raise NotImplementedError()

    timestamp = get_timestamp()
    order_id = order_id or generate_client_order_id()
    assert isinstance(order_id, int), \
        'order_id must be integer type: %s' % str(type(order_id))
    order = Order(order_id=order_id,
                  product=product,
                  price=price,
                  qty=qty,
                  order_side=order_side,
                  order_type=order_type,
                  tag=None,
                  accepted=False,
                  cancel_sent=False,
                  internal=BitfinexOrderInternal.create_by_submit(price_orig=price,
                                                                  qty_orig=qty,
                                                                  timestamp_created=timestamp))
    self._working_orders[order_id] = order
    self._og_logger.gen_order_request_and_order_event_and_log(order.product.symbol,
                                                              order.price,
                                                              order.qty,
                                                              order.order_side,
                                                              order.order_type,
                                                              str(order.order_id),
                                                              timestamp)
    self._og_logger.gen_order_response_and_order_event_and_log(order.product.symbol,
                                                               True,
                                                               str(order.order_id),
                                                               timestamp)
    return self._private_sub.place_order(product.native_symbol,
                                         None,
                                         str(price),
                                         post_qty,
                                         client_order_id=order.order_id,
                                         post_only=post_only,
                                         hidden=hidden,
                                         margin_trading=True)

  def cancel_multiple_impl(self, order_ids, *, ignore_error=False):
    if not isinstance(order_ids, list):
      order_ids = [order_ids]
    server_order_ids = []
    for order_id in order_ids:
      if order_id not in self._working_orders:
        self._logger.error('Cancel error: order %s not in working orders.' % order_id)
      try:
        order = self._working_orders[order_id]
        server_order_id = order.internal.server_order_id
        if server_order_id is None:
          raise KeyError()
        server_order_ids.append(server_order_id)
        # TODO(inkyu): Handle Failure
      except KeyError:
        self._logger.error('Cancel error: order %s is not accepted yet.' % order_id)
        continue

    return self._private_sub.cancel_by_server_order_ids(server_order_ids)

  def cancel_all_impl(self, *, ignore_error=False):
    order_ids = list(self._working_orders.keys())
    self.cancel_multiple(order_ids)

  def cancel_product_impl(self, product, *, ignore_error=False):
    order_ids = self.order_ids_by_product(product)
    if order_ids:
      self.cancel_multiple(order_ids)

  def order_ids_by_product(self, product):
    res = []
    for order in self._working_orders.values():
      if order.product == product:
        res.append(order.order_id)
    return res
